import cv2
import numpy as np

# 全局变量（替代类的实例属性）
img = None  # 原始图像
processed_img = None  # 处理后图像
img_path = None  # 图像路径
roi_img = None  # 选中的ROI区域


def menu():
    #显示主菜单
    while True:
        print(f"""
===== 图像处理工具箱 =====
    1. 加载图像
    2. 透视变换
    3. 图像旋转
    4. 对比度增强
    5. 亮度增强
    6. 添加水印
    7. 选择ROI区域
    8. 图像缩放
    9. 保存图像
    10. 保存ROI图像
    0. 退出
                    """)
        choice = input("请选择功能（输入数字）：")
        if choice == "1":
            load_image()
        elif choice == "2":
            perspective_image()
        elif choice == "3":
            rotate_image()
        elif choice == "4":
            enhance_contrast()
        elif choice == "5":
            enhance_brightness()
        elif choice == "6":
            add_watermark()
        elif choice == "7":
            select_roi()
        elif choice == "8":
            resize_image()
        elif choice == "9":
            save_image()
        elif choice == '10':
            save_roi_image()
        elif choice == "0":
            print("再见！")
            break
        else:
            print("无效输入，请重新选择！")


def load_image():
    #加载图像
    global img, processed_img, img_path
    path = input("请输入图像路径（如 D:/test.jpg）：").strip()
    img = cv2.imread(path)
    if img is None:
        print("图像加载失败！请检查路径是否正确或文件是否存在。")
        return
    processed_img = img.copy()  # 处理基于副本
    img_path = path
    print(f"图像加载成功！尺寸：{img.shape[0]}x{img.shape[1]}")


def perspective_image():
    #透视变换
    global img, processed_img
    if img is None:
        print("请先加载图像！")
        return

    points = []  # 存储选中的4个点
    window_name = "perspective"
    img_copy = processed_img.copy()  # 用于绘制的副本

    def mouse_callback(event, x, y, flags, param):
        nonlocal points, img_copy
        if event == cv2.EVENT_LBUTTONDOWN:
            if len(points) < 4:
                points.append((x, y))
                cv2.circle(img_copy, (x, y), 5, (0, 0, 255), -1)  # 红色标记点
                cv2.imshow(window_name, img_copy)

    # 显示图像并设置回调
    cv2.imshow(window_name, img_copy)
    cv2.setMouseCallback(window_name, mouse_callback)
    print("操作提示：左键点击选点（需按顺序选择4个角点），完成后按Enter确认，按ESC重新选择")

    # 等待用户确认
    while True:
        key = cv2.waitKey(0)
        if key == 13:  # Enter确认
            if len(points) == 4:
                cv2.destroyAllWindows()
                break
            else:
                print(f"需选择4个点，当前已选{len(points)}个，请继续选点")
        elif key == 27:  # ESC取消（清空当前选点）
            points = []
            img_copy = processed_img.copy()
            cv2.imshow(window_name, img_copy)
            print("已清空选点，请重新选择4个角点")

    # 执行透视变换
    pts1 = np.float32(points)
    h, w = processed_img.shape[:2]
    pts2 = np.float32([[0, 0], [w, 0], [w, h], [0, h]])  # 目标点
    M = cv2.getPerspectiveTransform(pts1, pts2)
    processed_img = cv2.warpPerspective(
        processed_img, M, (w, h),
        flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_WRAP
    )
    print("透视变换完成！")


def rotate_image():
    #图像旋转
    global img, processed_img
    if img is None:
        print("请先加载图像！")
        return
    angle = float(input("请输入旋转角度（如 90、-45）："))
    h, w = processed_img.shape[:2]
    center = (w // 2, h // 2)  # 图像中心坐标

    # 生成旋转矩阵并执行旋转
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    processed_img = cv2.warpAffine(processed_img, M, (w, h), flags=cv2.INTER_LINEAR)
    print(f"图像已旋转 {angle} 度！")


def draw_hist(gray_img, color):
    #绘制灰度图像的直方图
    # 计算直方图
    hist = cv2.calcHist([gray_img], [0], None, [256], [0, 256])
    # 归一化直方图（便于绘制）
    hist = cv2.normalize(hist, hist, 0, 255, cv2.NORM_MINMAX)
    # 创建直方图图像（高260，宽256）
    hist_img = np.ones((260, 256, 3), dtype=np.uint8) * 255  # 白色背景
    # 绘制直方图
    for i in range(256):
        cv2.line(
            hist_img,
            (i, 250),  # 起点（底部）
            (i, 250 - int(hist[i][0])),  # 终点（高度由直方图值决定）
            color,
            1
        )
    # 绘制坐标轴
    cv2.line(hist_img, (0, 250), (255, 250), (0, 0, 0), 2)  # x轴
    cv2.line(hist_img, (0, 0), (0, 250), (0, 0, 0), 2)  # y轴
    return hist_img


def enhance_contrast():
    #使用自适应直方图均衡化（CLAHE）增强对比度，并绘制直方图对比
    global img, processed_img
    if img is None:
        print("请先加载图像！")
        return

    # 1. 转换为灰度图（用于直方图计算）
    if len(processed_img.shape) == 3:
        gray_img = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY)
    else:
        gray_img = processed_img.copy()

    # 2. 绘制原始图像的直方图
    hist_original = draw_hist(gray_img, (0, 0, 255))  # 红色原始直方图

    # 3. 设置CLAHE参数（固定 clipLimit=3，tileGridSize=(80,80)）
    clip_limit = 3.0
    tile_size = 80
    print(f"使用固定参数：对比度限制={clip_limit}，分块大小={tile_size}x{tile_size}")

    # 4. 应用自适应直方图均衡化
    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(tile_size, tile_size))
    if len(processed_img.shape) == 3:
        # 彩色图：分通道处理
        b, g, r = cv2.split(processed_img)
        b_eq = clahe.apply(b)
        g_eq = clahe.apply(g)
        r_eq = clahe.apply(r)
        processed_img = cv2.merge([b_eq, g_eq, r_eq])
        # 均衡化后的灰度图（用于绘制直方图）
        gray_eq = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY)
    else:
        # 灰度图：直接处理
        processed_img = clahe.apply(gray_img)
        gray_eq = processed_img

    # 5. 绘制均衡化后的直方图
    hist_equalized = draw_hist(gray_eq, (0, 255, 255))  # 黄色均衡化直方图

    # 6. 显示直方图对比
    cv2.imshow("original_image", hist_original)
    cv2.imshow("processed_image", hist_equalized)
    print("自适应直方图均衡化完成！")
    print("按任意键关闭直方图窗口...")
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def enhance_brightness():
    #亮度增强
    global img, processed_img
    if img is None:
        print("请先加载图像！")
        return

    beta = int(input("请输入亮度增量（-100到100）："))
    if not (-100 <= beta <= 100):
        print("增量超出范围，使用默认值30")
        beta = 30

    # 亮度增强（对比度不变）
    processed_img = cv2.convertScaleAbs(processed_img, alpha=1, beta=beta)
    print("亮度增强完成！")


def add_watermark():
    #添加图像水印
    global img, processed_img
    if img is None:
        print("请先加载图像！")
        return

    # 输入水印图片路径
    watermark_path = input("请输入水印图像路径：").strip()
    watermark = cv2.imread(watermark_path)
    if watermark is None:
        print("水印图像加载失败！请检查路径是否正确。")
        return

    # 输入水印位置（左上角坐标）
    try:
        x = int(input("请输入水印左上角X坐标："))
        y = int(input("请输入水印左上角Y坐标："))
    except ValueError:
        print("坐标输入无效！")
        return

    # 水印预处理：灰度化→二值化→生成掩码
    gray_wm = cv2.cvtColor(watermark, cv2.COLOR_BGR2GRAY)
    _, mask_inv = cv2.threshold(gray_wm, 1, 255, cv2.THRESH_BINARY_INV)  # 反掩码
    mask = cv2.bitwise_not(mask_inv)  # 正掩码

    # 获取水印和原图尺寸
    h_watermark, w_watermark = watermark.shape[:2]
    h, w = processed_img.shape[:2]

    # 检查水印位置是否超出原图范围
    if x < 0 or y < 0 or x + w_watermark > w or y + h_watermark > h:
        print("水印位置超出原图范围，请重新选择坐标！")
        return

    # 提取原图中水印区域的ROI
    roi = processed_img[y:y + h_watermark, x:x + w_watermark].copy()

    # 位运算融合
    roi_background = cv2.bitwise_and(roi, roi, mask=mask_inv)
    wm_frontground = cv2.bitwise_and(watermark, watermark, mask=mask)
    dst = cv2.add(roi_background, wm_frontground)

    # 将融合结果放回原图
    processed_img[y:y + h_watermark, x:x + w_watermark] = dst
    print("图像水印添加完成！")


def select_roi():
    #选择ROI区域
    global img, processed_img, roi_img
    if img is None:
        print("请先加载图像！")
        return

    print("请在图像窗口中拖动鼠标选择ROI区域，按Enter确认，按c取消")
    roi = cv2.selectROI('choose_roi', processed_img, False)
    cv2.destroyAllWindows()  # 关闭ROI窗口

    if roi[2] > 0 and roi[3] > 0:  # 确保ROI有效
        x, y, w, h = roi
        roi_img = processed_img[y:y + h, x:x + w]  # 保存ROI
        print(f"ROI选择完成！尺寸：{h}x{w}")
    else:
        print("未选择有效ROI！")


def resize_image():
    #图像缩放
    global img, processed_img
    if img is None:
        print("请先加载图像！")
        return

    # 获取原图尺寸
    h, w = processed_img.shape[:2]
    print(f"当前图像尺寸：宽度 {w}px，高度 {h}px")

    # 输入X轴（宽度）和Y轴（高度）的缩放倍率
    try:
        scale_x = float(input("请输入X轴（宽度）缩放倍率（0.1-5.0）："))
        scale_y = float(input("请输入Y轴（高度）缩放倍率（0.1-5.0）："))
    except ValueError:
        print("输入无效！")
        return

    # 倍率范围限制
    scale_x = max(0.1, min(scale_x, 5.0))
    scale_y = max(0.1, min(scale_y, 5.0))

    # 计算新尺寸
    new_w = int(w * scale_x)
    new_h = int(h * scale_y)

    # 根据缩放方向选择插值方法
    if scale_x > 1 or scale_y > 1:
        interpolation = cv2.INTER_LINEAR  # 放大用双线性插值
    else:
        interpolation = cv2.INTER_AREA  # 缩小时用区域插值

    # 执行缩放
    processed_img = cv2.resize(processed_img, (new_w, new_h), interpolation=interpolation)
    print(f"图像缩放完成！")
    print(f"X轴倍率：{scale_x}，新宽度：{new_w}px")
    print(f"Y轴倍率：{scale_y}，新高度：{new_h}px")


def save_image():
    #保存处理后的完整图像
    global processed_img
    if processed_img is None:
        print("没有可保存的图像！")
        return

    save_path = input("请输入保存路径（如 D:/result.jpg）：").strip()
    # 自动补充扩展名
    if not (save_path.endswith(".jpg") or save_path.endswith(".png") or save_path.endswith(".bmp")):
        save_path += ".jpg"

    if cv2.imwrite(save_path, processed_img):
        print(f"图像已保存至：{save_path}")
    else:
        print("保存失败！请检查路径是否有效。")


def save_roi_image():
    #保存ROI区域的图像
    global roi_img
    if roi_img is None:
        print("没有可保存的ROI区域！请先选择ROI。")
        return

    save_path = input("请输入ROI区域的保存路径（如 D:/roi_result.jpg）：").strip()
    # 自动补充扩展名
    if not (save_path.endswith(".jpg") or save_path.endswith(".png") or save_path.endswith(".bmp")):
        save_path += ".jpg"

    if cv2.imwrite(save_path, roi_img):
        print(f"ROI区域已保存至：{save_path}")
    else:
        print("ROI保存失败！请检查路径是否有效。")


if __name__ == "__main__":
    # 确保中文显示正常
    menu()  # 启动菜单